from selenium import webdriver
from time import sleep, time
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
import requests
import request
import uuid
import fileUtils
import queryChromeHistory
from PIL import Image

import os
from selenium.webdriver.chrome.options import Options


def convert_webp_to_jpg(webp_path, jpg_path):
    try:
        # 打开WebP图像
        with Image.open(webp_path) as img:
            # 保存为JPEG格式
            img.convert("RGB").save(jpg_path, "JPEG")
        print(f"成功将 {webp_path} 转换为 {jpg_path}")
    except Exception as e:
        print(f"转换失败: {e}")


def startFetch(url, tabIndex=0):
    """开始抓取脚本"""
    # os.system("taskkill /f /im chrome.exe")

    # 设置下载路径
    download_path = r"H:\download\tmp"
    chrome_history_path = r"H:\download\chrome\profile"
    # 设置 ChromeOptions 来配置下载路径和自动下载
    chrome_options = Options()
    chrome_options.add_experimental_option("prefs", {
        "download.default_directory": download_path,
        "download.prompt_for_download": False,
        "download.directory_upgrade": True,
        "safebrowsing.enabled": True,
        "profile.default_content_settings.popups": 0
    })

    # 指定用户数据目录
    chrome_options.add_argument(r"--user-data-dir=H:\download\userdata")
    # 下载文件数据库的保存地址
    chrome_options.add_argument(r"--user-data-dir=" + chrome_history_path)
    chrome_options.add_argument("--start-minimized")  # 设置窗口默认最小化

    # 1.创建Chrome浏览器对象，这会在电脑上在打开一个浏览器窗口
    service = Service('D:\chromedriver-win64\chromedriver.exe')
    driver = webdriver.Chrome(service=service, options=chrome_options)

    # 4.设置浏览器的大小
    driver.set_window_size(1500, 700)
    # 2.通过浏览器向服务器发送URL请求
    driver.get(url)
    sleep(1)
    # 判断是否登录
    while True:
        findWxLogin = False
        findQQLogin = False
        try:
            loginElement = driver.find_element(By.XPATH, "//*[contains(text(), '登录/注册')]")
            findWxLogin = True
            sleep(1)
        except:
            sleep(1)

        try:
            loginElement = driver.find_element(By.XPATH, "//*[contains(text(), '使用你的QQ头像')]")
            findQQLogin = True
            print("未登录,等待登录后再试")
            sleep(1)
        except:
            findQQLogin = False
            sleep(1)

        if findWxLogin is not True and findQQLogin is not True:
            print("已登录")
            break
        else:
            print("未登录,等待登录后再试")

    # 判断当前选择的tab
    # 查找标题
    tabElement = driver.find_element(By.CSS_SELECTOR,
                                     "div.ant-tabs-nav-wrap div.ant-tabs-nav-list")
    print(tabElement.text)
    tabElements = tabElement.find_elements(By.XPATH, './div')
    # 切换到选择的tab
    if tabIndex < len(tabElements):
        tabElements[tabIndex].click()

    # 使用CSS选择器找到h1元素
    # 查找标题
    h1_element = driver.find_element(By.CSS_SELECTOR, "div.ModelInfoHead_firstLeft__LuHui span")
    # 获取h1元素的文本值
    nameOrigin = h1_element.text

    # 遍历表格根据名称查询
    xiangqingElement = driver.find_element(By.XPATH, "//*[contains(text(), '版本详情')]")
    # t1 = xiangqingElement.text
    parentElement = xiangqingElement.find_element(By.XPATH, "./..")  # 使用 "./.." 表示父级
    # t2 = parentElement.text
    tableElement = parentElement.find_element(By.CSS_SELECTOR, 'div:nth-child(2)')
    tableElement = tableElement.find_elements(By.XPATH, './div')
    # 在div元素下查找所有的a标签
    type = ""
    version = ""
    triggerWord = ""
    for element in tableElement:
        titleElement = element.find_element(By.CSS_SELECTOR, 'div')

        # 查找lora和ckpt类型
        if "类型" in titleElement.text:
            tdElement2 = element.find_element(By.CSS_SELECTOR, 'div:nth-child(2)')
            type = tdElement2.text
        # 查找基础模型
        if "基础算法" in titleElement.text:
            tdElement2 = element.find_element(By.CSS_SELECTOR, 'div:nth-child(2)')
            version = tdElement2.text
        # 查找关键词
        if "触发词" in titleElement.text:
            tdElement2 = element.find_element(By.CSS_SELECTOR,
                                              'div:nth-child(2)')
            wordsElement = tdElement2.find_elements(By.XPATH, './div')
            for e in wordsElement:
                triggerWord = triggerWord + e.text + ","

    # 查找首图
    imgElement = driver.find_element(By.CSS_SELECTOR,
                                     "div.flex-shrink-0 div.items-center div.relative img")
    imgUrl = imgElement.get_attribute('src')

    # 查找下载地址
    downloadUrl = ""
    canDownload = True
    try:
        downloadElement = driver.find_element(By.XPATH, "//*[contains(text(), '不可下载')]")
        print(downloadElement.text)
        canDownload = False
        downloadUrl = "不可下载"
    except:
        canDownload = True

    if canDownload:
        try:
            print("开始抓取下载地址的方式")
            # 首先查找点击显示下载按钮的方式
            # xiazaiElement = driver.find_element(By.XPATH, "//*[contains(text(), '下载')]")
            xiazaiElement = driver.find_elements(By.CSS_SELECTOR,
                                                 "i.icon-border_download")

            # 检查是否有至少两个匹配的元素
            if len(xiazaiElement) >= 2:
                xiazaiElement = xiazaiElement[1]
                print(xiazaiElement.text)
                # 找到该元素的父元素
                downloadElement = xiazaiElement.find_element(By.XPATH, "./..")  # 使用 "./.." 表示父级

                downloadElement.click()
                # 获取重定向后的URL
                # 使用JavaScript模拟点击下载按钮
                sleep(2)

            # 判断是否有弹窗
            downloadElement = driver.find_element(By.XPATH,
                                                  "//*[contains(text(), '直接下载加密模型')]")
            downloadElement.click()
            # 获取重定向后的URL
            # 使用JavaScript模拟点击下载按钮
            sleep(2)
        except:
            print("地址抓取失败")

    driver.quit()  # 退出整个WebDriver会话，关闭所有窗口
    sleep(1)
    if canDownload:
        localFilePath, url = queryChromeHistory.getLatestDownload(
            chrome_history_path=chrome_history_path + r"\Default\History")
        downloadUrl = url

    # 下载图片
    # 生成一个UUID
    new_uuid = uuid.uuid4()
    # 将UUID转换为字符串
    uuid_str = str(new_uuid)
    suffix = ""
    if ",webp" in imgUrl:
        suffix = "webp"
    else:
        suffix = fileUtils.getFileNameSuffix(imgUrl)

    localImgPath = download_path + "\\" + uuid_str + "." + suffix
    request.download_file(imgUrl, localImgPath)

    if ",webp" in imgUrl:
        newLocalImgPath = download_path + "\\" + uuid_str + ".jpg"
        # 把图片转换为jpg
        convert_webp_to_jpg(localImgPath, newLocalImgPath)
        # 删除源文件
        os.remove(localImgPath)
        localImgPath = newLocalImgPath

    # 打印文本值
    print("头图地址", imgUrl)
    print("头图本地地址", localImgPath)
    print("标题为:", nameOrigin)
    print("类型为:", type)
    print("基础模型:", version)
    print("关键词:", triggerWord)
    print("最终下载地址为：", downloadUrl)
    # sleep(1)

    return imgUrl, localImgPath, nameOrigin, type, version, triggerWord, downloadUrl


if __name__ == '__main__':
    # startFetch("https://civitai.com/models/84728/photon")
    # startFetch("https://www.liblib.art/modelinfo/a8ca174c51e641eba2b28e617a888980")
    # startFetch("https://www.liblib.art/modelinfo/b19ef6df34ac4972a109617c1840e008")
    # startFetch("https://www.liblib.art/modelinfo/b986ea62d51049d6b7d382b06b8e067a")  # 不可下载的
    # startFetch("https://www.liblib.art/modelinfo/f2af8733021a4b3e84592559e326b631")
    # startFetch(
    #     "https://www.liblib.art/modelinfo/770608ce91f54627a31c8ddf1a6d5b49?from=feed")  # 现代水墨
    # startFetch(
    #     "https://www.liblib.art/modelinfo/eb54f3880dd04cc4b9002d31753a12ec?from=feed")  # 儿童简笔
    # startFetch("https://www.liblib.art/modelinfo/9be966833d164528add75ee989522866",
    #            tabIndex=1)  # 多选项卡
    # startFetch("https://www.liblib.art/modelinfo/ea92e254b8904f34b06332b762fb22b9",
    #            tabIndex=0)  # 多选项卡
    startFetch(
        "https://www.liblib.art/modelinfo/204ae3dedee348b4b2ac9d6ee273bd51?from=feed&versionUuid=de3597abdfda4473991669ee48e9c7b6",
        tabIndex=0)  # 多选项卡
    # startFetch("https://www.liblib.art/modelinfo/1e7fcce2a48b4916ab97ffac3aa6e64d?from=feed&versionUuid=22a8b16421724a79976bdf461234714d",
    #            tabIndex=0)  # 多选项卡
